import numpy as np
import os
import cv2
import datetime
import time
import re
import shutil

videoPath = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/round2.avi'
savePath = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/anchorImgs_method2'

anchorImgPath = os.path.join(savePath, 'anchorImgs') # 保存的是锚点图像
datasetPath = os.path.join(savePath, 'dataset') # 保存的是包括及其相邻帧的图像，用于训练

if os.path.exists(datasetPath):
    ans = input('Dataset path %s already exit, sure to remove and create a new one? [y]/n : '%datasetPath)
    if ans == 'n':
        print('Program exit.')
        exit(0)
    shutil.rmtree(datasetPath)
imgNames = []
for _,_, imgNames in os.walk(anchorImgPath):
    pass
imgNames = [i[0:10] for i in imgNames]
imgNamesInt = list(map(int, imgNames))

sampleNum = 16 
sample_interval = 1

sampleImgInt = []
last_positive_idx = [-1]
min_positive_idx_different = np.inf
for i,_ in enumerate(imgNamesInt):
    if i == len(imgNamesInt) - 1:
        break
    anchor_idx = imgNamesInt[i]
    positive_idx = [anchor_idx-i*sample_interval for i in range(int(sampleNum/2))] + [anchor_idx+(i+1)*sample_interval for i in range(int(sampleNum/2))]
    positive_idx.sort()

    cur_positive_idx_different = min(positive_idx) - max(last_positive_idx)
    if cur_positive_idx_different <= 0:
        print('Error! Positive samples of different anchor mixed! Program exit.')
        exit(0)
    if cur_positive_idx_different < min_positive_idx_different:
        min_positive_idx_different = cur_positive_idx_different

    sampleImgInt.append(positive_idx)
    last_positive_idx = positive_idx

print('Min positive samples index difference between neighbor anchors is %d'%min_positive_idx_different)

cap = cv2.VideoCapture(videoPath)
for oneAnchorImgs in sampleImgInt:
    folderPath = str(oneAnchorImgs[0]).rjust(10, '0')
    folderPath = os.path.join(datasetPath, folderPath)
    os.makedirs(folderPath)
    for imgFno in oneAnchorImgs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, imgFno)
        fno = cap.get(cv2.CAP_PROP_POS_FRAMES)
        if fno != imgFno:
            raise ValueError('Frame number error!!! It should be %d, but it is %d'%(imgFno, fno))
        imgPath = str(imgFno).rjust(10,'0') + '.png'
        imgPath = os.path.join(folderPath, imgPath)
        ret, img = cap.read()
        if not ret:
            raise ValueError('Cannot read frame %d!'%imgFno)
        cv2.imwrite(imgPath, img)